package com.yxh.www.gateway.filter;

import com.alibaba.fastjson.JSONObject;
import com.yxh.www.common.result.ResultBuilder;
import com.yxh.www.common.result.ResultEnum;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * <p>
 *  令牌桶实现 限流拦截器
 * </p>
 *
 * @author yangxiaohui
 * @since 2020/5/14
 */
@Component
public class RateLimitByIpGatewayFilter implements GlobalFilter, Ordered {
    int capacity;//桶的最大容量，即能装载 令牌 的最大数量
    int refillTokens;//每次 令牌 补充量
    Duration refillDuration;//补充 令牌 的时间间隔
    private static final Map<String, Bucket> CACHE = new ConcurrentHashMap<>();
    private static final boolean enableRateLimit = true;//开关
    public RateLimitByIpGatewayFilter() {
        this.capacity=50;
        this.refillTokens=5;
        this.refillDuration=Duration.ofSeconds(1);
    }

    private Bucket createNewBucket() {
        Refill refill = Refill.of(refillTokens, refillDuration);
        Bandwidth limit = Bandwidth.classic(capacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        if (!enableRateLimit){
            return chain.filter(exchange);
        }
        String ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
        Bucket bucket = CACHE.computeIfAbsent(ip, k -> createNewBucket());
        if (bucket.tryConsume(1)) {
            return chain.filter(exchange);
        } else {
            return this.buildRateLimitMessage(exchange.getResponse());
        }
    }

    private Mono<Void> buildRateLimitMessage(ServerHttpResponse resp){
        resp.setStatusCode(HttpStatus.OK);
        resp.getHeaders().add("Content-Type", "application/json;charset=UTF-8");
        // 构造错误信息
        String returnStr = JSONObject.toJSONString(ResultBuilder.error(ResultEnum.TOO_MANY_REQUESTS, "请求过于频繁！"));
        DataBuffer buffer = resp.bufferFactory().wrap(returnStr.getBytes(StandardCharsets.UTF_8));
        return resp.writeWith(Flux.just(buffer));
    }


    @Override
    public int getOrder() {
        // 第一个执行这个过滤器
        return Ordered.HIGHEST_PRECEDENCE;
    }
}
